#pragma GCC optimize(2)
#pragma GCC optimize(3,"Ofast","inline")
#include "kernel_operator.h"

using namespace AscendC;

#define ull uint64_t
// #define ll int64_t

template<typename T> class KernelDepthToSpace {
private:
    GlobalTensor<T> xGm;
    GlobalTensor<T> yGm;
    int8_t mode, data_format;
    ull block_size, n, h, w, c, cbb;
    ull shape[6];
    ull trans[6];
    ull orig[6];
    ull suf[7];
    ull transsuf[7];

public:
    __aicore__ inline KernelDepthToSpace() {}
    __aicore__ inline void Go(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, ull totalLength, ull tileNum, int8_t mode, int8_t data_format, ull block_size, ull n, ull h, ull w, ull c, ull cbb, ull input_size)
    {
        xGm.SetGlobalBuffer((__gm__ T *)x, totalLength);
        yGm.SetGlobalBuffer((__gm__ T *)y, totalLength);

        this->mode = mode;
        this->data_format = data_format;
        this->block_size = block_size;
        this->n = n;
        this->h = h;
        this->w = w;
        this->c = c;
        this->cbb = cbb;
        
        if (mode == 0 && data_format == 0) {
            // DCR, NHWC
            ull shape0[6] = {n, h, w, block_size, block_size, cbb};
            ull trans0[6] = {0, 1, 3, 2, 4, 5};
            for(int i=0;i<6;i++) shape[i]=shape0[i],trans[i]=trans0[i];
        } else if (mode == 0 && data_format == 1) {
            // DCR, NCHW
            ull shape0[6] = {n, block_size, block_size, cbb, h, w};
            ull trans0[6] = {0, 3, 4, 1, 5, 2};
            for(int i=0;i<6;i++) shape[i]=shape0[i],trans[i]=trans0[i];
        } else if (mode == 1 && data_format == 0) {
            // CRD, NHWC
            ull shape0[6] = {n, h, w, cbb, block_size, block_size};
            ull trans0[6] = {0, 1, 4, 2, 5, 3};
            for(int i=0;i<6;i++) shape[i]=shape0[i],trans[i]=trans0[i];
        } else if (mode == 1 && data_format == 1) {
            // CRD, NCHW
            ull shape0[6] = {n, cbb, block_size, block_size, h, w};
            ull trans0[6] = {0, 1, 4, 2, 5, 3};
            for(int i=0;i<6;i++) shape[i]=shape0[i],trans[i]=trans0[i];
        } else {
            assert(0, "Error: mode or data_format is invalid\n");
        }

        for (int i = 0; i < 6; i++) {
            orig[trans[i]] = i;
        }
        
        suf[6] = transsuf[6] = 1;
        for (int i = 5; i >= 0; i--) {
            suf[i] = shape[i];
            transsuf[i] = shape[trans[i]];
            suf[i] *= suf[i + 1];
            transsuf[i] *= transsuf[i + 1];
        }

        for (ull i0 = 0; i0 < shape[0]; i0++) {
            for (ull i1 = 0; i1 < shape[1]; i1++) {
                for (ull i2 = 0; i2 < shape[2]; i2++) {
                    for (ull i3 = 0; i3 < shape[3]; i3++) {
                        for (ull i4 = 0; i4 < shape[4]; i4++) {
                            for (ull i5 = 0; i5 < shape[5]; i5++) {
                                ull i[6] = {i0, i1, i2, i3, i4, i5};
                                ull input_index = 0;
                                ull output_index = 0;
                                for (int j = 0; j < 6; j++) {
                                    input_index += i[j] * suf[j+1];
                                    output_index += i[trans[j]] * transsuf[j+1];
                                }
                                T tmp = xGm.GetValue(input_index);
                                yGm.SetValue(output_index, tmp);
                            }
                        }
                    }
                }
            }
        }
    }
};

constexpr ull BUFFER_NUM = 16;
constexpr ull TQueDepth = 16;

template<typename T> class KernelDepthToSpace00 {
private:
    TPipe* pipe;
    GlobalTensor<T> xGm;
    GlobalTensor<T> yGm;
    TQueBind<TPosition::VECIN, TPosition::VECOUT, TQueDepth> inQueueX;

public:
    __aicore__ inline KernelDepthToSpace00() {}
    __aicore__ inline void Go(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, ull blockLength, ull tileLength, uint32_t tileSize, ull shape2, ull shape3, ull suf3, ull suf4, ull transsuf3, ull transsuf4, uint32_t bufferSize, uint32_t stride)
    {
        /*
        blockLength = totalLength / GetBlockNum()
        tileLength = cb
        tileSize = cb * input_size
        shape2 = w
        shape3 = block_size
        suf3 = c
        suf4 = cb
        transsuf3 = cb*w
        transsuf4 = cb
        */
        pipe = pipeIn;

        ull offset = blockLength * GetBlockIdx();
        xGm.SetGlobalBuffer((__gm__ T *)x + offset, blockLength);
        yGm.SetGlobalBuffer((__gm__ T *)y + offset, blockLength);

        LocalTensor<T> xLocal;
        
        pipe->InitBuffer(inQueueX, BUFFER_NUM, bufferSize);
        
        // 输入顺序
        // ull xidx = 0;
        // ull yidx0 = 0;
        // ull yidx;
        // for (ull i = 0; i < shape2; i++) {
        //     yidx = yidx0;
        //     for (ull j = 0; j < shape3; j++) {
        //             xLocal = inQueueX.AllocTensor<T>();
        //             DataCopy(xLocal, xGm[xidx], tileLength);
        //             inQueueX.EnQue(xLocal);
        //             xLocal = inQueueX.DeQue<T>();
        //             DataCopy(yGm[yidx], xLocal, tileLength);
        //             inQueueX.FreeTensor(xLocal);
        //         xidx += suf4;
        //         yidx += transsuf3;
        //     }
        //     yidx0 += transsuf4;
        // }

        // 输入顺序，去掉第二层循环
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParamsIn{1, bufferSize, 0, 0, 0};
        DataCopyExtParams copyParamsOut{(uint16_t)shape3, tileSize, 0, stride, 0};
        ull xidx = 0;
        ull yidx = 0;
        for (ull i = 0; i < shape2; i++) {
                xLocal = inQueueX.AllocTensor<T>();
                DataCopyPad(xLocal, xGm[xidx], copyParamsIn, padParams);
                inQueueX.EnQue(xLocal);
                xLocal = inQueueX.DeQue<T>();
                DataCopyPad(yGm[yidx], xLocal, copyParamsOut);
                inQueueX.FreeTensor(xLocal);
            xidx += suf3;
            yidx += transsuf4;
        }

        // 输出顺序，去掉第二层循环
        // uint32_t stride = suf3 * sizeof(T) - tileSize;
        // DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        // DataCopyExtParams copyParamsIn{(uint16_t)shape2, tileSize, 0, stride, 0};
        // DataCopyExtParams copyParamsOut{1, bufferSize, 0, 0, 0};
        // ull xidx = 0;
        // ull yidx = 0;
        // for (ull i = 0; i < shape3; i++) {
        //         xLocal = inQueueX.AllocTensor<T>();
        //         DataCopyPad(xLocal, xGm[xidx], copyParamsIn, padParams);
        //         inQueueX.EnQue(xLocal);
        //         xLocal = inQueueX.DeQue<T>();
        //         DataCopyPad(yGm[yidx], xLocal, copyParamsOut);
        //         inQueueX.FreeTensor(xLocal);
        //     xidx += suf4;
        //     yidx += transsuf3;
        // }

        // 输出顺序
        // ull yidx = 0;
        // ull xidx0 = 0;
        // ull xidx;
        // for (ull j = 0; j < shape3; j++) {
        //     xidx = xidx0;
        //     for (ull i = 0; i < shape2; i++) {
        //             xLocal = inQueueX.AllocTensor<T>();
        //             DataCopy(xLocal, xGm[xidx], tileLength);
        //             inQueueX.EnQue(xLocal);
        //             xLocal = inQueueX.DeQue<T>();
        //             DataCopy(yGm[yidx], xLocal, tileLength);
        //             inQueueX.FreeTensor(xLocal);
        //         xidx += suf3;
        //         yidx += transsuf4;
        //     }
        //     xidx0 += suf4;
        // }
    }
};

extern "C" __global__ __aicore__ void depth_to_space(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    int8_t mode = tiling_data.mode;
    int8_t data_format = tiling_data.data_format;
    if (TILING_KEY_IS(1)) {
        KernelDepthToSpace<DTYPE_X> op;
        op.Go(x, y, workspace, &pipe, tiling_data.totalLength, tiling_data.tileNum, tiling_data.mode, tiling_data.data_format, tiling_data.block_size, tiling_data.n, tiling_data.h, tiling_data.w, tiling_data.c, tiling_data.cbb, tiling_data.input_size);
    } else if (TILING_KEY_IS(2)){
        KernelDepthToSpace00<DTYPE_X> op;
        op.Go(x, y, workspace, &pipe, tiling_data.blockLength, tiling_data.tileLength, tiling_data.tileSize, tiling_data.shape2, tiling_data.shape3, tiling_data.suf3, tiling_data.suf4, tiling_data.transsuf3, tiling_data.transsuf4, tiling_data.bufferSize, tiling_data.stride);
    }
}